from langgraph.graph import END, StateGraph
from typing import Dict, TypedDict, List, Any
import re
import json

# 定义工具结构
class Tool(TypedDict):
    name: str
    description: str
    parameters: Dict[str, Any]

# 定义大模型响应结构
class ModelResponse(TypedDict):
    thought: str
    tool: str
    parameters: Dict[str, Any]

# 定义状态结构
class AgentState(TypedDict):
    user_input: str          # 用户输入的字符串
    model_response: ModelResponse  # 大模型结构化响应
    tool_result: Any         # 工具执行结果
    final_output: str        # 最终输出给用户的内容

# 模拟可用工具
TOOLS: List[Tool] = [
    {
        "name": "add",
        "description": "执行两个数字的加法运算",
        "parameters": {
            "type": "object",
            "properties": {
                "a": {"type": "number", "description": "第一个数字"},
                "b": {"type": "number", "description": "第二个数字"}
            },
            "required": ["a", "b"]
        }
    },
    {
        "name": "subtract",
        "description": "执行两个数字的减法运算",
        "parameters": {
            "type": "object",
            "properties": {
                "a": {"type": "number", "description": "被减数"},
                "b": {"type": "number", "description": "减数"}
            },
            "required": ["a", "b"]
        }
    }
]

# 模拟大模型接口 - 解析用户输入并选择工具
def mock_llm_parse_input(state: AgentState) -> Dict:
    """模拟大模型解析用户输入并选择工具"""
    print(f"\n🧠 用户输入: '{state['user_input']}'")
    print(f"🛠️ 可用工具: {[t['name'] for t in TOOLS]}")

    # 使用正则表达式提取数字
    numbers = re.findall(r'-?\d+', state["user_input"])

    # 简单规则判断操作类型
    if "加" in state["user_input"] or "+" in state["user_input"] or "和" in state["user_input"]:
        tool_name = "add"
        thought = "用户要求进行加法运算"
    elif "减" in state["user_input"] or "-" in state["user_input"] or "差" in state["user_input"]:
        tool_name = "subtract"
        thought = "用户要求进行减法运算"
    else:
        # 默认使用加法
        tool_name = "add"
        thought = "未明确指定操作类型，默认使用加法"

    # 确保提取到两个数字
    if len(numbers) < 2:
        # 模拟大模型尝试修复
        thought += " - 输入中数字不足，补充默认值0"
        numbers += ["0", "0"]  # 补充默认值
        numbers = numbers[:2]  # 只取前两个

    # 转换为整数
    try:
        a = int(numbers[0])
        b = int(numbers[1])
    except ValueError:
        a, b = 0, 0
        thought += " - 数字解析失败，使用默认值0"

    # 构建结构化响应
    response: ModelResponse = {
        "thought": thought,
        "tool": tool_name,
        "parameters": {"a": a, "b": b}
    }

    print(f"🤖 模拟大模型响应: {json.dumps(response, indent=2, ensure_ascii=False)}")
    return {"model_response": response}

# 工具执行函数
def execute_tool(state: AgentState) -> Dict:
    """根据模型选择的工具执行相应操作"""
    tool_name = state["model_response"]["tool"]
    params = state["model_response"]["parameters"]
    a = params["a"]
    b = params["b"]

    print(f"🔧 执行工具: {tool_name}(a={a}, b={b})")

    if tool_name == "add":
        result = a + b
        print(f"➕ 加法结果: {a} + {b} = {result}")
    elif tool_name == "subtract":
        result = a - b
        print(f"➖ 减法结果: {a} - {b} = {result}")
    else:
        result = None
        print(f"⚠️ 未知工具: {tool_name}")

    return {"tool_result": result}

# 生成最终响应
def generate_final_response(state: AgentState) -> Dict:
    """生成最终输出给用户的内容"""
    tool_name = state["model_response"]["tool"]
    params = state["model_response"]["parameters"]
    result = state["tool_result"]
    thought = state["model_response"]["thought"]

    if result is None:
        output = "抱歉，我无法完成这个计算"
    elif tool_name == "add":
        output = f"{params['a']} + {params['b']} = {result}"
    elif tool_name == "subtract":
        output = f"{params['a']} - {params['b']} = {result}"
    else:
        output = f"计算结果: {result}"

    # 添加思考过程
    output = f"{thought}\n\n{output}"

    print(f"📝 最终输出: {output}")
    return {"final_output": output}

# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("parse_input", mock_llm_parse_input)
    workflow.add_node("execute_tool", execute_tool)
    workflow.add_node("generate_response", generate_final_response)

    # 设置入口点
    workflow.set_entry_point("parse_input")

    # 连接节点（这里不需要条件路由，因为所有路径都一样）
    workflow.add_edge("parse_input", "execute_tool")
    workflow.add_edge("execute_tool", "generate_response")
    workflow.add_edge("generate_response", END)

    return workflow.compile()

# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 测试案例
    test_cases = [
        {"input": "5加3等于多少？", "expected": "5 + 3 = 8"},
        {"input": "10减去7", "expected": "10 - 7 = 3"},
        {"input": "计算20和15的和", "expected": "20 + 15 = 35"},
        {"input": "100 - 50", "expected": "100 - 50 = 50"},
        {"input": "8和4的差是多少？", "expected": "8 - 4 = 4"},
        {"input": "简单的计算：3和5", "expected": "3 + 5 = 8"},
        {"input": "只有数字：42 100", "expected": "42 + 100 = 142"},
        {"input": "计算圆的面积", "expected": "抱歉，我无法完成这个计算"},
    ]

    for test in test_cases:
        print("\n" + "="*50)
        print(f"测试输入: '{test['input']}'")
        result = app.invoke({"user_input": test["input"]})

        # 验证结果
        final_output = result["final_output"]
        if test["expected"] in final_output:
            print(f"✅ 测试通过! 输出包含: '{test['expected']}'")
        else:
            print(f"❌ 测试失败! 期望: '{test['expected']}'，实际: '{final_output}'")

        print(f"完整状态: {json.dumps(result, indent=2, ensure_ascii=False)}")
